use reqwest::header;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::Value;

#[derive(Debug)]
pub struct TyClient {
    api_key: String,
    client: Client,
}

#[derive(Deserialize, Debug, Serialize)]
pub struct MemoryData {
    user: String,
    robot: String,
}
#[derive(Serialize, Deserialize)]
struct InputMessage {
    role: String,
    content: String,
}

#[derive(Serialize, Deserialize)]
struct ChatModel {
    model: String,
    input: InputData,
    parameters: serde_json::Value,
}

#[derive(Serialize, Deserialize)]
struct InputData {
    messages: Vec<InputMessage>,
}
impl TyClient {
    // 新建一个客户端，并指定api
    pub fn new(api_key: String) -> TyClient {
        let client = Client::new();
        TyClient { api_key, client }
    }

    pub async fn handle_sse(
        &self,
        user_msg: &str,
        memory: Vec<MemoryData>,
    ) -> Result<reqwest::Response, reqwest::Error> {
        let mut messages = vec![InputMessage {
            role: "system".to_string(),
            content: "You are a helpful assistant.".to_string(),
        }];

        for data in &memory {
            messages.push(InputMessage {
                role: "user".to_string(),
                content: data.user.clone(),
            });
            messages.push(InputMessage {
                role: "assistant".to_string(),
                content: data.robot.clone(),
            });
        }

        messages.push(InputMessage {
            role: "user".to_string(),
            content: user_msg.to_string(),
        });

        // Build the ChatModel
        let chat_model = ChatModel {
            model: "qwen-72b-chat".to_string(),
            input: InputData { messages },
            parameters: serde_json::Value::Object(serde_json::Map::new()),
        };

        // Serialize the ChatModel to JSON
        let body = serde_json::to_value(&chat_model).expect("Failed to serialize to JSON");

        Ok(self.post_sse(body).await?)
    }

    pub async fn post_sse(&self, body: Value) -> Result<reqwest::Response, reqwest::Error> {
        // 请求头
        let mut headers = header::HeaderMap::new();
        headers.insert(
            header::AUTHORIZATION,
            format!("Bearer {}", self.api_key).parse().unwrap(),
        );
        headers.insert(header::CONTENT_TYPE, "application/json".parse().unwrap());
        headers.insert("X-DashScope-SSE", "enable".parse().unwrap());

        // 请求体
        let request_body = serde_json::to_string(&body).unwrap();
        let response = self
            .client
            .post("https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation")
            .headers(headers)
            .body(request_body)
            .send()
            .await?;
        Ok(response)
    }
}
